/**
 * Copyright (c) iwindplus Technologies Co., Ltd.2024-2030, All rights reserved.
 */

package com.iwindplus.boot.web;

import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import com.alibaba.ttl.TtlRunnable;
import com.alibaba.ttl.threadpool.TtlExecutors;
import com.iwindplus.boot.web.domain.property.ThreadPoolProperty;
import com.iwindplus.boot.web.thread.TtlThreadPoolTaskExecutor;
import com.iwindplus.boot.web.thread.TtlThreadPoolTaskScheduler;
import jakarta.annotation.Resource;
import java.time.Duration;
import java.util.Arrays;
import java.util.Optional;
import java.util.concurrent.Executor;
import java.util.concurrent.ThreadPoolExecutor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.aop.interceptor.AsyncUncaughtExceptionHandler;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.autoconfigure.task.TaskExecutionProperties;
import org.springframework.boot.autoconfigure.task.TaskSchedulingProperties;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.scheduling.annotation.AsyncConfigurer;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;

/**
 * 线程池配置.
 *
 * @author zengdegui
 * @since 2023/08/29 22:24
 */
@Slf4j
@EnableAsync
@EnableScheduling
@Configuration(proxyBeanMethods = false)
@EnableConfigurationProperties(ThreadPoolProperty.class)
@ConditionalOnProperty(prefix = "spring.task", name = "enabled", havingValue = "true", matchIfMissing = true)
public class ThreadPoolConfiguration implements AsyncConfigurer {

    @Resource
    private TaskExecutionProperties taskExecutionProperties;

    @Resource
    private TaskSchedulingProperties taskSchedulingProperties;

    @Override
    public Executor getAsyncExecutor() {
        final String simpleName = ThreadPoolTaskExecutor.class.getSimpleName();
        final String beanName = StrUtil.lowerFirst(simpleName);
        return SpringUtil.getBean(beanName, ThreadPoolTaskExecutor.class);
    }

    /**
     * 创建 ThreadPoolTaskExecutor.
     *
     * @return ThreadPoolTaskExecutor
     */
    @Primary
    @Bean
    public ThreadPoolTaskExecutor threadPoolTaskExecutor() {
        TtlThreadPoolTaskExecutor executor = new TtlThreadPoolTaskExecutor();
        final TaskExecutionProperties.Pool pool = this.taskExecutionProperties.getPool();
        final TaskExecutionProperties.Shutdown shutdown = this.taskExecutionProperties.getShutdown();
        executor.setCorePoolSize(Optional.ofNullable(pool.getCoreSize()).orElse(Runtime.getRuntime().availableProcessors()));
        executor.setMaxPoolSize(Optional.ofNullable(pool.getMaxSize()).orElse(Runtime.getRuntime().availableProcessors() * 2));
        executor.setQueueCapacity(Optional.ofNullable(pool.getQueueCapacity()).orElse(executor.getQueueCapacity()));
        executor.setKeepAliveSeconds((int) Optional.ofNullable(pool.getKeepAlive()).orElse(Duration.ofSeconds(60)).getSeconds());
        executor.setAllowCoreThreadTimeOut(Optional.ofNullable(pool.isAllowCoreThreadTimeout()).orElse(Boolean.TRUE));
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.AbortPolicy());
        executor.setThreadNamePrefix(this.taskExecutionProperties.getThreadNamePrefix());
        executor.setWaitForTasksToCompleteOnShutdown(Optional.ofNullable(shutdown.isAwaitTermination()).orElse(Boolean.TRUE));
        executor.setAwaitTerminationSeconds(
            (int) Optional.ofNullable(shutdown.getAwaitTerminationPeriod()).orElse(Duration.ofSeconds(60)).getSeconds());
        executor.setTaskDecorator(TtlRunnable::get);
        executor.initialize();
        TtlExecutors.getTtlExecutor(executor);
        log.info("ThreadPoolTaskExecutor={}", executor);
        return executor;
    }

    @Override
    public AsyncUncaughtExceptionHandler getAsyncUncaughtExceptionHandler() {
        return (ex, method, params) ->
            log.warn("[ThreadPool Exception] class={}, method={}, params={}, error={}",
                method.getDeclaringClass().getSimpleName(), method.getName(), Arrays.toString(params), ex);
    }


    /**
     * 创建 ThreadPoolTaskScheduler.
     *
     * @return ThreadPoolTaskScheduler
     */
    @Primary
    @Bean
    @ConditionalOnProperty(prefix = "spring.task", name = "enabled-scheduler", havingValue = "true", matchIfMissing = true)
    public ThreadPoolTaskScheduler threadPoolTaskScheduler() {
        TtlThreadPoolTaskScheduler scheduler = new TtlThreadPoolTaskScheduler();
        final TaskSchedulingProperties.Pool pool = this.taskSchedulingProperties.getPool();
        final TaskSchedulingProperties.Shutdown shutdown = this.taskSchedulingProperties.getShutdown();
        scheduler.setPoolSize(Optional.ofNullable(pool.getSize()).orElse(Runtime.getRuntime().availableProcessors()));
        scheduler.setRejectedExecutionHandler(new ThreadPoolExecutor.AbortPolicy());
        scheduler.setThreadNamePrefix(this.taskSchedulingProperties.getThreadNamePrefix());
        scheduler.setAwaitTerminationSeconds(
            (int) Optional.ofNullable(shutdown.getAwaitTerminationPeriod()).orElse(Duration.ofSeconds(60)).getSeconds());
        scheduler.setWaitForTasksToCompleteOnShutdown(Optional.ofNullable(shutdown.isAwaitTermination()).orElse(Boolean.TRUE));
        scheduler.initialize();
        TtlExecutors.getTtlExecutor(scheduler);
        log.info("ThreadPoolTaskScheduler={}", scheduler);
        return scheduler;
    }
}